#include "StdAfx.h"
#include "PEFile.h"

static const int MAX_FILE_SIZE = 30*1024*1024;
static const int BUFFER_BYTES_EACH_TIME = 1024 * 1024;

CPEFile::CPEFile(void):
m_pDosHeader(NULL),
m_pNtHeader(NULL),
m_pFileHeader(NULL),
m_pOptionalHeader(NULL),
m_pSectionHeader(NULL),
m_nSectionsNum(0)
{
	m_nFileSize = 0;
	m_pvMapAddr = NULL;
	m_pvMapAddr = VirtualAlloc(NULL, MAX_FILE_SIZE, MEM_COMMIT, PAGE_READWRITE);
}

CPEFile::~CPEFile(void)
{
	VirtualFree(m_pvMapAddr, MAX_FILE_SIZE, MEM_RELEASE);
}

BOOL CPEFile::Open(const CString & strFileName)
{
	HANDLE hFile = CreateFile(strFileName, GENERIC_READ, 
		FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_FLAG_SEQUENTIAL_SCAN, NULL);

	if (hFile == INVALID_HANDLE_VALUE)
		return FALSE;

	DWORD dwReadBytes = 0;
	char * buff = (char *) m_pvMapAddr;
	m_nFileSize = 0;

	while(ReadFile(hFile, buff, BUFFER_BYTES_EACH_TIME, &dwReadBytes, NULL))
	{
		buff += BUFFER_BYTES_EACH_TIME;
		m_nFileSize += dwReadBytes;
	}

	CloseHandle(hFile);

	return TRUE;
}


BOOL CPEFile::IsPEFile()
{
	PIMAGE_DOS_HEADER  pDH=NULL;
	PIMAGE_NT_HEADERS  pNtH=NULL;

	if(!m_pvMapAddr)
		return FALSE;

	pDH=(PIMAGE_DOS_HEADER)m_pvMapAddr;
	if(pDH->e_magic!=IMAGE_DOS_SIGNATURE)
		return FALSE;

	pNtH=(PIMAGE_NT_HEADERS32)((DWORD)pDH+pDH->e_lfanew);
	if (pNtH->Signature != IMAGE_NT_SIGNATURE )
		return FALSE;

	return TRUE;
}

BOOL CPEFile::ParsePE()
{
	ASSERT(m_pvMapAddr != NULL);

	if (IsPEFile() == FALSE)
		return FALSE;

	m_pDosHeader = (PIMAGE_DOS_HEADER)m_pvMapAddr;
	m_pNtHeader = (PIMAGE_NT_HEADERS)((DWORD)m_pDosHeader+m_pDosHeader->e_lfanew);
	m_pFileHeader = &m_pNtHeader->FileHeader;
	m_pOptionalHeader = &m_pNtHeader->OptionalHeader;
	m_nSectionsNum = m_pFileHeader->NumberOfSections;
	m_pSectionHeader = IMAGE_FIRST_SECTION(m_pNtHeader);

	return TRUE;
}


BOOL CPEFile::AddSection(LPCSTR strSectionName, LPVOID pvContent, 
				DWORD pvVirtualAddr, DWORD dwVirtualSize, 
				DWORD dwRawDataSize, DWORD dwRawDataOffset, 
				DWORD dwSectionFlag)
{
	IMAGE_SECTION_HEADER sectionHeader;
	memset(&sectionHeader, 0, sizeof(sectionHeader));
	strncpy((char*)&sectionHeader.Name[0], strSectionName, sizeof(sectionHeader.Name));
	sectionHeader.Misc.VirtualSize = dwVirtualSize;
	sectionHeader.VirtualAddress = (DWORD)pvVirtualAddr;
	sectionHeader.SizeOfRawData = dwRawDataSize;
	sectionHeader.PointerToRawData = dwRawDataOffset;
	sectionHeader.PointerToLinenumbers = 0;
	sectionHeader.PointerToRelocations = 0;
	sectionHeader.NumberOfLinenumbers = 0;
	sectionHeader.NumberOfRelocations = 0;
	sectionHeader.Characteristics = dwSectionFlag;

	PIMAGE_SECTION_HEADER pNewSectionHeader= &m_pSectionHeader[m_pFileHeader->NumberOfSections];
	memcpy(pNewSectionHeader, &sectionHeader, sizeof(IMAGE_SECTION_HEADER));
	m_pOptionalHeader->SizeOfImage += dwVirtualSize;
	m_pOptionalHeader->SizeOfCode += dwVirtualSize;
	m_pFileHeader->NumberOfSections++;
	memcpy(&((char *)m_pvMapAddr)[dwRawDataOffset], pvContent, dwRawDataSize);

	m_nFileSize = dwRawDataSize + dwRawDataOffset;

	return TRUE;
}

BOOL CPEFile::Save(const CString & strFileName)
{
	HANDLE hFile = CreateFile(strFileName, GENERIC_READ | GENERIC_WRITE, 
		FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, FILE_FLAG_SEQUENTIAL_SCAN, NULL);

	if (hFile == INVALID_HANDLE_VALUE)
		return FALSE;

	char * buff = (char *)m_pvMapAddr;
	DWORD dwBytesWritten = 0;

	for (int i = 0; i < m_nFileSize; i += BUFFER_BYTES_EACH_TIME)
	{
		int nBytesToWrite = BUFFER_BYTES_EACH_TIME;
		if ((m_nFileSize - i) < BUFFER_BYTES_EACH_TIME)
		{
			nBytesToWrite = (m_nFileSize - i);
		}
		if (FALSE == WriteFile(hFile, buff, nBytesToWrite, &dwBytesWritten, NULL))
			return FALSE;
	}

	CloseHandle(hFile);

	return TRUE;
}
